import uuid
import asyncio
from dataclasses import dataclass
from datetime import datetime, timedelta
from io import BytesIO
from math import ceil

from django.conf import settings
from django.utils.timezone import now

import structlog
from google.genai import Client
from google.genai.errors import APIError
from google.genai.types import Blob, Content, Part, VideoMetadata
from pymediainfo import MediaInfo
from temporalio.common import RetryPolicy, WorkflowIDReusePolicy

from posthog.models.exported_asset import ExportedAsset
from posthog.models.user import User
from posthog.settings.temporal import TEMPORAL_WORKFLOW_MAX_ATTEMPTS
from posthog.storage import object_storage
from posthog.sync import database_sync_to_async
from posthog.temporal.common.client import async_connect
from posthog.temporal.exports_video.workflow import VideoExportInputs, VideoExportWorkflow

from products.llm_analytics.backend.providers.gemini import GeminiProvider

from ee.hogai.session_summaries.constants import (
    DEFAULT_VIDEO_EXPORT_MIME_TYPE,
    DEFAULT_VIDEO_UNDERSTANDING_MODEL,
    SHORT_VALIDATION_VIDEO_PLAYBACK_SPEED,
    VALIDATION_VIDEO_DURATION,
    VALIDATION_VIDEO_RENDERING_DELAY,
)

logger = structlog.get_logger(__name__)


@dataclass(frozen=True)
class SessionMomentInput:
    # ID to identify the moment in mappings (for example, event_uuid)
    moment_id: str
    # Timestamp to start the video from
    timestamp_s: int
    # How long the video should be
    duration_s: int
    # Prompt to validate the moment
    prompt: str


@dataclass(frozen=True)
class SessionMomentOutput(SessionMomentInput):
    # Asset ID of the stored video
    asset_id: int
    # Description of the moment, generated by LLM
    video_description: str
    # When the video was created
    created_at: datetime
    # When the video will expire
    expires_after: datetime
    # What model was used to analyze the video
    model_id: str


class SessionMomentsLLMAnalyzer:
    """Generate videos for Replay session events and analyze them with LLM"""

    def __init__(self, session_id: str, team_id: int, user: User, trace_id: str | None = None):
        self.session_id = session_id
        self.team_id = team_id
        self.user = user
        self.trace_id = trace_id
        self._provider = GeminiVideoUnderstandingProvider(model_id=DEFAULT_VIDEO_UNDERSTANDING_MODEL)

    async def analyze(
        self, moments_input: list[SessionMomentInput], expires_after_days: int, failed_moments_min_ratio: float
    ) -> list[SessionMomentOutput]:
        """
        Analyze the session moments with LLM and return mapping of moment_id to LLM analysis

        Args:
            moments_input: List of session moments to analyze
            expires_after_days: For how long to store the videos. Depends on the use case.
            failed_moments_min_ratio: If less than N% of moments were generated, fail the analysis. Depends on the use case.
        """
        # Generate mapping of moments to moment ids
        moment_id_to_moment = {moment.moment_id: moment for moment in moments_input}
        # Generate mapping of created videos (asset IDs) to moment ids
        moment_id_to_asset_id = await self._generate_videos_for_moments(
            moments_input=moments_input,
            expires_after_days=expires_after_days,
            failed_moments_min_ratio=failed_moments_min_ratio,
        )
        # Analyze videos with LLM
        results = await self._analyze_moment_videos_with_llm(
            moment_id_to_asset_id=moment_id_to_asset_id,
            moment_id_to_moment=moment_id_to_moment,
            expires_after_days=expires_after_days,
        )
        return results

    async def _generate_videos_for_moments(
        self,
        moments_input: list[SessionMomentInput],
        expires_after_days: int,
        failed_moments_min_ratio: float,
    ) -> dict[str, int]:
        """Generate videos for moments and return mapping of moment_id to asset_id"""
        tasks = {}
        async with asyncio.TaskGroup() as tg:
            for moment in moments_input:
                tasks[moment.moment_id] = tg.create_task(
                    self._generate_video_for_single_moment(moment=moment, expires_after_days=expires_after_days)
                )
        # Collect asset IDs
        moment_to_asset_id: dict[str, int] = {}
        for moment_id, task in tasks.items():
            res: int | Exception = task.result()
            if isinstance(res, Exception):
                logger.exception(
                    f"Failed to generate video for moment {moment_id} from session {self.session_id} of team {self.team_id}: {res}"
                )
                # Not failing explicitly to avoid failing all the generations if one fails
                continue
            moment_to_asset_id[moment_id] = res
        # Check if enough moments were generated
        expected_min_moments = ceil(len(moments_input) * failed_moments_min_ratio)
        if expected_min_moments > len(moment_to_asset_id):
            exception_message = f"Not enough moments were generated for session {self.session_id} of team {self.team_id}: {len(moment_to_asset_id)} out of {len(moments_input)}, expected at least {expected_min_moments}"
            logger.exception(exception_message)
            # Remove all the generated videos to not bloat the database
            asset_ids = list(moment_to_asset_id.values())
            await ExportedAsset.objects.filter(id__in=asset_ids).adelete()
            raise Exception(exception_message)
        return moment_to_asset_id

    def _generate_moment_video_filename(self, moment_id: str) -> str:
        """Generate a filename for a moment video"""
        return f"session-moment_{self.session_id}_{moment_id}"

    async def _generate_video_for_single_moment(
        self, moment: SessionMomentInput, expires_after_days: int
    ) -> int | Exception:
        """Generate a video for an event in Replay session and return the asset ID"""
        try:
            moment_filename = self._generate_moment_video_filename(moment_id=moment.moment_id)
            created_at = now()
            expires_after = created_at + timedelta(days=expires_after_days)
            exported_asset = await ExportedAsset.objects.acreate(
                team_id=self.team_id,
                export_format=DEFAULT_VIDEO_EXPORT_MIME_TYPE,
                export_context={
                    "session_recording_id": self.session_id,
                    "timestamp": moment.timestamp_s,
                    "filename": moment_filename,
                    "duration": moment.duration_s,
                    # Speed up to reduce the rendering time
                    "playback_speed": SHORT_VALIDATION_VIDEO_PLAYBACK_SPEED,
                    # Keeping default values
                    "mode": "screenshot",
                },
                created_by=self.user,
                created_at=created_at,
                expires_after=expires_after,
            )
            # Generate a video through Temporal workflow
            client = await async_connect()
            await client.execute_workflow(
                VideoExportWorkflow.run,
                VideoExportInputs(exported_asset_id=exported_asset.id),
                id=f"session-moment-video-export_{self.session_id}_{moment.moment_id}_{uuid.uuid4()}",
                task_queue=settings.VIDEO_EXPORT_TASK_QUEUE,
                retry_policy=RetryPolicy(maximum_attempts=int(TEMPORAL_WORKFLOW_MAX_ATTEMPTS)),
                id_reuse_policy=WorkflowIDReusePolicy.ALLOW_DUPLICATE_FAILED_ONLY,
            )
            # Return the asset ID for later retrieval
            return exported_asset.id
        except Exception as err:  # Workflow retries exhausted
            # Let caller handle the error
            return err

    async def _get_video_bytes(self, asset_id: int) -> bytes | None:
        """Retrieve video content as bytes for an ExportedAsset ID"""
        try:
            content: bytes | None = None
            # Fetch the asset from the database
            asset = await ExportedAsset.objects.aget(id=asset_id)
            # Get content from either database or object storage
            if asset.content:
                # Content stored directly in database
                content = bytes(asset.content)
            elif asset.content_location:
                # Content stored in object storage
                content = await database_sync_to_async(object_storage.read_bytes, thread_sensitive=False)(
                    asset.content_location
                )
            if not content:
                logger.exception(
                    f"No content found for asset {asset_id} for session {self.session_id} "
                    f"of team {self.team_id} when validating session summaries"
                )
                raise ExportedAsset.DoesNotExist
            return content
        except ExportedAsset.DoesNotExist:
            return None

    @staticmethod
    def _get_webm_duration(video_bytes: bytes) -> int | None:
        """Extract duration in milliseconds from WEBM video bytes to understand when the export UI finished rendering"""
        try:
            media_info = MediaInfo.parse(BytesIO(video_bytes))
            for track in media_info.tracks:
                if track.track_type == "General":
                    # Convert ms to seconds, ceil to avoid grey "not-rendered" frames at the start
                    return ceil(track.duration / 1000.0)
            return None
        except Exception as e:
            logger.exception(f"Error extracting video duration: {e}")
            return None

    async def _analyze_single_moment_video_with_llm(
        self,
        asset_id: int,
        moment_id: str,
        prompt: str,
    ) -> str | None | Exception:
        """Analyze a moment video with LLM"""
        try:
            # Not chunking bytes, as even with 300 sessions at once it's 100MB memory tops
            video_bytes = await self._get_video_bytes(asset_id=asset_id)
            if not video_bytes:
                logger.warning(
                    f"No video bytes found for asset {asset_id} for moment {moment_id} of session {self.session_id} of team {self.team_id}"
                )
                return None
            # Calculate how many seconds to skip from the start, as Puppeteer needs time to render the export UI
            total_video_duration = self._get_webm_duration(video_bytes=video_bytes)
            start_offset_s = (
                total_video_duration - VALIDATION_VIDEO_DURATION + VALIDATION_VIDEO_RENDERING_DELAY
                if total_video_duration
                else None
            )
            # Analyze the video with LLM
            content = await self._provider.understand_video(
                video_bytes=video_bytes,
                mime_type=DEFAULT_VIDEO_EXPORT_MIME_TYPE,
                prompt=prompt,
                start_offset_s=start_offset_s,
                # No end offset is required, as we expect the export rendering to stop right after the video was played
                trace_id=self.trace_id,
            )
            if not content:
                logger.warning(
                    f"No LLM content found for moment {moment_id} of session {self.session_id} of team {self.team_id}"
                )
                return None
            return content
        except Exception as err:
            logger.exception(
                f"Failed to analyze moment video {moment_id} of session {self.session_id} of team {self.team_id} with LLM: {err}"
            )
            # If the LLM validation fails - ensure to remove the generated video to not bloat the database,
            # as it would be linked to the summary (to reuse) only after the LLM video validation is completed
            await ExportedAsset.objects.filter(id=asset_id).adelete()
            return err  # Let caller handle the error

    async def _analyze_moment_videos_with_llm(
        self,
        moment_id_to_asset_id: dict[str, int],
        moment_id_to_moment: dict[str, SessionMomentInput],
        expires_after_days: int,
    ) -> list[SessionMomentOutput]:
        """Send videos to LLM for validation and get analysis results"""
        tasks = {}
        async with asyncio.TaskGroup() as tg:
            for moment_id, asset_id in moment_id_to_asset_id.items():
                prompt = moment_id_to_moment[moment_id].prompt
                tasks[moment_id] = tg.create_task(
                    self._analyze_single_moment_video_with_llm(asset_id=asset_id, moment_id=moment_id, prompt=prompt)
                )
        results: list[SessionMomentOutput] = []
        for moment_id, task in tasks.items():
            res = task.result()
            if isinstance(res, Exception):
                logger.exception(
                    f"Failed to analyze moment video {moment_id} of session {self.session_id} of team {self.team_id} with LLM: {res}"
                )
                continue
            if not res:
                continue
            moment = moment_id_to_moment[moment_id]
            created_at = now()
            expires_after = created_at + timedelta(days=expires_after_days)
            output = SessionMomentOutput(
                moment_id=moment.moment_id,
                timestamp_s=moment.timestamp_s,
                duration_s=moment.duration_s,
                prompt=moment.prompt,
                asset_id=moment_id_to_asset_id[moment_id],
                video_description=res,
                created_at=created_at,
                expires_after=expires_after,
                model_id=self._provider.model_id,
            )
            results.append(output)
        # No additional check for how many moments were analyzed as they can be limited by video size
        return results


class GeminiVideoUnderstandingProvider:
    """Interface for Gemini video understanding"""

    # https://ai.google.dev/gemini-api/docs/video-understanding#supported-formats
    SUPPORTED_VIDEO_MIME_TYPES: list[str] = [
        "video/x-flv",
        "video/quicktime",
        "video/mpeg",
        "video/mpegs",
        "video/mpg",
        "video/mp4",
        "video/webm",
        "video/wmv",
        "video/3gpp",
    ]

    VIDEO_MAX_SIZE_BYTES = 20 * 1024 * 1024  # 20MB

    def __init__(self, model_id: str):
        self.model_id = model_id
        # Using PostHog Gemini provider to avoid logic duplication
        self._base_provider = GeminiProvider(model_id=model_id)
        # Using default Gemini client as workaround, as PostHog wrapper doesn't support async yet
        self.client = Client(api_key=self._base_provider.get_api_key())

    async def understand_video(
        self,
        video_bytes: bytes,
        mime_type: str,
        prompt: str,
        start_offset_s: int | None = None,
        end_offset_s: int | None = None,
        trace_id: str | None = None,
    ) -> str | None:
        """
        Understand a video and return a summary using the provided prompt
        https://ai.google.dev/gemini-api/docs/video-understanding
        """
        self._base_provider.validate_model(self.model_id)
        if mime_type not in self.SUPPORTED_VIDEO_MIME_TYPES:
            logger.exception(
                f"Video bytes for understanding video are not in a supported MIME type (trace_id:{trace_id}): {mime_type}"
            )
            return None
        if not len(video_bytes):
            logger.exception(f"Video bytes for understanding video are empty (trace_id: {trace_id})")
            return None
        if len(video_bytes) > self.VIDEO_MAX_SIZE_BYTES:
            logger.exception(f"Video bytes for understanding video are too large (trace_id: {trace_id})")
            return None
        try:
            video_part_config: dict[str, str | Blob | VideoMetadata] = {
                "inline_data": Blob(data=video_bytes, mime_type=mime_type)
            }
            video_metadata_config = {}
            if start_offset_s:
                video_metadata_config["start_offset"] = f"{start_offset_s}s"
            if end_offset_s:
                video_metadata_config["end_offset"] = f"{end_offset_s}s"
            if video_metadata_config:
                video_part_config["video_metadata"] = VideoMetadata(**video_metadata_config)
            video_part = Part(**video_part_config)
            prompt_part = Part(text=prompt)
            contents = Content(parts=[video_part, prompt_part])
            response = await self.client.aio.models.generate_content(
                model=self.model_id,
                contents=contents,
                # Add trace ID, when PostHog wrapper supports async
            )
            return response.text
        except APIError as e:
            logger.exception(f"Gemini API error while understanding video (trace_id: {trace_id}): {e}")
            return None
        except Exception as e:
            logger.exception(f"Unexpected error while understanding video (trace_id: {trace_id}): {e}")
            return None
